Repository navigation
Update the longrope inv_freq buffer in place - #50480
Conversation
`longrope_frequency_update` re-registered `inv_freq` on every forward, replacing the buffer with a new tensor. CUDA graphs replay from the buffer's original address, so the replaced tensor is never read and vLLM rejects the module under CUDA graphs. Copy into the existing buffer instead. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
[For maintainers] Suggested jobs to run (before merge) run-slow: phi3 |
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
zucchini-nlp
left a comment
zucchini-nlp
left a comment
There was a problem hiding this comment.
looks reasonable to me, thx!
|
run-slow: phi3 |
| inv_freq = torch.where(seq_len > original_max_position_embeddings, long_inv_freq, original_inv_freq) | ||
| self.register_buffer(f"{prefix}inv_freq", inv_freq, persistent=False) | ||
| # Update in place so the buffer keeps its address, which CUDA graphs replay from | ||
| getattr(self, f"{prefix}inv_freq").copy_(inv_freq) |
There was a problem hiding this comment.
wondering whether we need something similar for dynamic freq updates?
There was a problem hiding this comment.
it is still not compile friendly overall (input-dependent seq_len > max_len, jiqing tried but it led to computational overhead at each forward step
There was a problem hiding this comment.
This was motivated by vLLM, but ultimately I now must choose either long or short depending on how vLLM is configured. This is how vLLM treats log rope for its native models too
CI (Nvidia + AMD)Nvidia · ✅ No new failures 🎉👏 · ⚙️ Workflow · 📊 Dashboard Failures (1 pre-existing)Pre-existing failures (still failing) — 1 total: Same error (1): - phi3: AMD · ✅ No new failures 🎉👏 · ⚙️ Workflow · 📊 Dashboard Failures (5 pre-existing)Pre-existing failures (still failing) — 5 total: Same error (5): - phi3: |
What does this PR do?
Follow-up to #49555.
longrope_frequency_updatenow picks the long or short frequencies withtorch.where, but it still re-registersinv_freqon every forward, replacing the buffer with a new tensor. Under CUDA graphs the captured kernels keep reading the original buffer, so the switch never reaches them; vLLM rejects the module for this reason. This PR copies into the existing buffer instead, so it keeps its address.Evidence
test_longrope_updates_inv_freq_in_placechecks the buffer keeps its identity while switching long → short; it fails onmain.test_longrope_compiled_matches_eager(from #49555) still passes.End to end,
microsoft/Phi-3.5-mini-instructthrough vLLM's Transformers modeling backend (B200, bf16,max_model_len=8192, piecewise CUDA graphs), compared with vLLM's native Phi3:mainAssigning / modifying buffers of nn.Module during forward pass is not allowed when using cudagraphThe vLLM side also needs vllm-project/vllm#65446, which narrows that check to real
_bufferswrites.AI assistance (Claude Code) was used to write this PR; I reviewed and validated it.
Before submitting
Pull Request checks?
to it if that's the case.
Who can review?
@vasqu @zucchini-nlp (reviewed #49555)